#!/usr/bin/env python3

import rclpy
import cv2
import numpy as np
import manipulation.scripts.constants.detection as const
from cv_bridge import CvBridge
from manipulation.scripts.detection.commans.base_detection_node import BaseDetectionNode
from manipulation.scripts.detection.apriltag.apriltag_detection import AprilTagDetection
from scipy.spatial.transform import Rotation
from threading import Thread

class ApriltagDetectionNode(BaseDetectionNode):
    def __init__(self):
        super().__init__('apriltag_detection_node')
        self.declare_parameter('tag_size', 0.04)
        self.tag_size = self.get_parameter('tag_size').get_parameter_value().double_value

        self.detect_cfg = const.apriltag_tf_publish_config
        self.april_tag_detection = AprilTagDetection()
        self.cv_bridge = CvBridge()
        self.thread = Thread(target=self.detect, daemon=True)
        self.thread.start()

        self.get_logger().info(f"ApriltagDetectionNode initialized with tag size: {self.tag_size}, detection config: {self.detect_cfg}")

    def detect(self):
        while True:
            rgb_img, depth_img = self.get_frame()
            img, tag_infos = self.april_tag_detection.draw_detected_corners(
                rgb_img=rgb_img,
                K=self.camera.K,
                dist=np.array(self.camera.dist, np.float32),
                draw_text=False,
                tag_size=self.tag_size
            )

            self.publish_ideal_camera_tf()
            targets_list = []
            for cfg in self.detect_cfg:
                tag_idx = cfg['tag_idx']
                child_frame_id = cfg['child_frame_id']
                if tag_idx in tag_infos:
                    infos = tag_infos[tag_idx]
                    if len(infos) > 1:
                        self.get_logger().warn(f"More than one tag found for {tag_idx}, using the first one.")
                    if tag_idx in const.tag_index_targets:
                        targets_list.append({
                            'child_frame_id': child_frame_id,
                            'R': infos[0]['R'],
                            't': infos[0]['t'],
                        })
                    self.publish_tf_from_detect_info(child_frame_id, infos[0])
                    img = self.draw_axis_in_image(img, infos[0]['R'], infos[0]['t'], child_frame_id)
            self.targets_list = targets_list

            # Publish image
            img_msg = self.cv_bridge.cv2_to_imgmsg(img, encoding='bgr8')
            img_msg.header.stamp = self.get_clock().now().to_msg()
            img_msg.header.frame_id = 'camera'
            self.pub_detected_img.publish(img_msg)


if __name__ == '__main__':
    rclpy.init()
    node = ApriltagDetectionNode()
    rclpy.spin(node)
    node.destroy_node()
    rclpy.shutdown()
